import torch
from resnet_impl import resnet18
import onnxruntime as ort
import numpy as np
import cv2
import time


class Resnet(object):
    def __init__(self):
        weights = 'C:/Users/liyu/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth'
        self.net = resnet18(weights)
        self.net.eval()
    
    def infer(self, inputs):
        return self.net(inputs)

    def export(self, onnx_path, input_size=224):
        inputs = torch.randn(1, 3, input_size, input_size)
        torch.onnx.export(self.net, inputs, onnx_path, 
            verbose=True,
            opset_version=11, 
            input_names=['input'], 
            output_names=['output'], 
        )


class SampleOnnx(object):
    def __init__(self, onnxpath):
        self.session = ort.InferenceSession(onnxpath)
    
    def infer(self, x):
        input_dict = {
            'input': x,
        }
        output_names = ['output']
        outputs = self.session.run(output_names, input_dict)
        return outputs[0]


def test_onnx(onnxpath, input_size):
    inputs = np.ones((1, 3, input_size, input_size), dtype=np.float32)
    demo = SampleOnnx(onnxpath)
    st = time.time()
    y = demo.infer(inputs)
    time.sleep(0.1)
    et = time.time()
    print(st)
    print(et)
    print('infer time: ', et - st)
    # print(y)
    print(y.shape)
    print(np.sum(y))


def load_image(image_path):
    image = cv2.imread(image_path)
    image = image[:, :, ::-1]
    image = cv2.resize(image, (224, 224))
    scale = 1.0 / 255.0
    mean = [0.5, 0.5, 0.5]
    std = [0.5, 0.5, 0.5]
    
    image = (image * scale - mean) / std
    image = image.astype(np.float32)
    image = image.transpose((2, 0, 1))
    return image


if __name__ == '__main__':
    onnx_path = '../data/resnet.onnx'
    input_size = 224
    # if True:
    #     model = Resnet()
    #     model.export(onnx_path, input_size)
    test_onnx(onnx_path, input_size)

    # image = load_image('../data/anime_sample.jpg')
    # print(image)
    # print(image.shape)
